
import torch
class ConfusionMatrix:
    def __init__(self, classes_num=2) -> None:
        
        self.classes_num = classes_num
        self.confusion_matrix = torch.zeros((classes_num, classes_num), requires_grad=False).long()

    def update_matrix(self, pred_y, true_y):
        for i in range(self.classes_num):
            for j in range(self.classes_num):
                self.confusion_matrix[i][j] = torch.sum((pred_y == i) & (true_y == j))
        return None

from torch import Tensor
class Metrics:
    def __init__(self, class_num=2, metric_type="macro") -> None:
        self.class_num = class_num
        self.metric_type= metric_type
        pass
    
    def __call__(self, pred_y:Tensor, true_y:Tensor) -> dict:
        if self.metric_type == "macro":
            return self.macro_metrics(pred_y, true_y)
        elif self.metric_type == "micro":
            return self.micro_metrics(pred_y, true_y)
        else:
            raise

    def micro_metrics(self, pred_y:Tensor, true_y:Tensor):
        class_num = self.class_num
        
        total_tp = 0
        total_fp = 0
        total_fn = 0
        total_tn = 0
        
        for i in range(class_num):
            tp_i = torch.sum((pred_y == i) & (true_y == i)).item()
            fp_i = torch.sum((pred_y == i) & (true_y != i)).item()
            fn_i = torch.sum((pred_y != i) & (true_y == i)).item()
            tn_i = torch.sum((pred_y != i) & (true_y != i)).item()
            
            total_tp += tp_i
            total_fp += fp_i
            total_fn += fn_i
            total_tn += tn_i
        
        mic_precision = total_tp/(total_tp+total_fp)
        mic_recall = total_tp/(total_tp+total_fn)
        mic_accuracy = (total_tp+total_tn)/(total_tp+total_fp+total_tn+total_fn)
        mic_f1 = 2*total_tp / (2*total_tp + total_fp + total_fn)
        
        return {
            "precision":mic_precision,
            "recall":mic_recall,
            "accuracy":mic_accuracy,
            "f1":mic_f1
        }
    
    def macro_metrics(self, pred_y:Tensor, true_y:Tensor):
        class_num = self.class_num
        total_precision = 0.0
        total_recall = 0.0
        total_accuracy = 0.0  # torch.sum(pred_y == true_y) / true_y.numel()
        total_f1 = 0.0
        
        for i in range(class_num):
            tp_i = torch.sum((pred_y == i) & (true_y == i)).item()
            fp_i = torch.sum((pred_y == i) & (true_y != i)).item()
            fn_i = torch.sum((pred_y != i) & (true_y == i)).item()
            tn_i = torch.sum((pred_y != i) & (true_y != i)).item()
            
            precision_i = tp_i/(tp_i + fp_i + 1)
            recall_i = tp_i/(tp_i + fn_i + 1)
            accuracy_i = (tp_i + tn_i)/(tp_i + fp_i + fn_i + tn_i)
            f1_i = 2*tp_i / (2*tp_i + fp_i + fn_i + 1)
            
            total_precision += precision_i
            total_recall += recall_i
            total_accuracy += accuracy_i
            total_f1 += f1_i
        
        mac_precision = total_precision / class_num
        mac_recall = total_recall / class_num
        mac_accuracy = total_accuracy / class_num
        mac_f1 = total_f1 / class_num
        
        return {
            "precision":mac_precision,
            "recall":mac_recall,
            "accuracy":mac_accuracy,
            "f1":mac_f1
        }
